//
//  unionFind.swift
//  others
//
//  Created by lhc on 2024/8/19.
//

import Foundation

protocol UnionFind {
    func getSize() -> Int
    func isConnected(p: Int, q: Int) -> Bool
    func unionElement(p: Int, q: Int)
}


class UnionFindLevel1: UnionFind {
    
    fileprivate var parents: [Int]
    
    init(size: Int) {
        self.parents = [Int]()
        for i in 0..<size {
            self.parents.append(i)
        }
    }
    
    func getSize() -> Int {
        return self.parents.count
    }
    
    fileprivate func findRoot(index: Int) -> Int {
        // 寻找到index在parent里指向的值
        return self.parents[index]
    }
    
    func isConnected(p: Int, q: Int) -> Bool {
        guard checkIndex(p: p, q: q) else { return false }
        // 判断index为p和q在parent里保存的指向的值是否相同
        return self.findRoot(index: p) == self.findRoot(index: q)
    }
    
    func unionElement(p: Int, q: Int) {
        guard checkIndex(p: p, q: q) else { return }
        
        guard self.findRoot(index: p) != self.findRoot(index: q) else { return }
        
        // 保存p的指向的值
        let target = self.findRoot(index: p)
        for i in 0..<self.getSize() {
            // 如果值等于p指向的值
            if self.findRoot(index: i) == target {
                // 设置成指向q
                self.parents[i] = self.findRoot(index: q)
            }
        }
    }
    
    
    fileprivate func checkIndex(p: Int, q: Int) -> Bool {
        if  p < 0
            || q < 0
            || p >= self.getSize()
            || q >= self.getSize()
        {
            return false
        }else {
            return true
        }
    }
}

class UnionFindLevel2: UnionFindLevel1 {
    // 修改为只要当前index在parent指向的不是自己，就继续查找直到找到指向自己的index的root
    override func findRoot(index: Int) -> Int {
        var root = self.parents[index]
        while root != self.parents[root] {
            root = self.parents[root]
        }
        
        return root
    }
    
    override func unionElement(p: Int, q: Int) {
        let rootP = self.findRoot(index: p)
        let rootQ = self.findRoot(index: q)
        
        guard rootP != rootQ else { return }
        
        // 让p的根index指向的index 指向q的根index指向的index
        self.parents[rootP] = rootQ
    }
}

class UnionFindLevel3: UnionFindLevel2 {
    
    
    var rank = [Int]()
    override init(size: Int) {
        super.init(size: size)
        // 保存每个根节点的rank值 就是这个树的高度
        self.rank = Array(repeating: 1, count: self.getSize())
    }
    
    override func unionElement(p: Int, q: Int) {
        guard self.checkIndex(p: p, q: q) else { return }
        
        let rootP = self.findRoot(index: p)
        let rootQ = self.findRoot(index: q)
        
        guard rootP != rootQ else { return }
        
        // 让高度低的树去指向高度高的树，接起来的树高度为高度高的树
        if self.rank[rootP] < self.rank[rootQ] {
            self.parents[rootP] = rootQ
        }else if rank[rootP] > self.rank[rootQ] {
            self.parents[rootQ] = rootP
        }else {
            // 这里两个树高度相等, 拼接的话rootQ的树高度会加1
            self.parents[rootP] = rootQ
            self.rank[rootQ] += 1
        }
    }
}

class UnionFindLevel4: UnionFindLevel2 {
    // 递归查找root的index时同时把index都指向root的index
    override func findRoot(index: Int) -> Int {
        
        if index != self.parents[index] {
            self.parents[index] = self.findRoot(index: self.parents[index])
        }
        return self.parents[index]
    }
}
